open Ast
open Printf
open Map
  
module NameMap =
  Map.Make
    (struct type t = string

             let compare x y = Pervasives.compare x y
                end)
  
exception ReturnException of int * int NameMap.t
  
(* These are the symbol tables *)
(* global symbol table [var name, var type]*)
let globals = NameMap.empty
  
(* local symbol table [var name, var type]*)
let locals = NameMap.empty
  
(* function declartion symbol table [func name, func decl]*)
let func_decls = NameMap.empty
  
(* returns a formatted name without $ *)
let get_name name = String.sub name 1 ((String.length name) - 1)

(* is the name given a keyword? *)
let is_keyword name =
  match name with
  | "print" -> true
  | "int" -> true
  | "string" -> true
  | "if" -> true
  | "else" -> true
  | "return" -> true
  | "for" -> true
  | "break" -> true
  | "while" -> true
  | "true" -> true
  | "false" -> true
  | "goto" -> true
  | "bool" -> true
	| "value" -> true
	| "deck" -> true
	| "getName" -> true
	| "getValue" -> true
  | _ -> false
  
(* Get datatype of value*)
let get_data_type locals globals func_decls =
  function
  | IntLiteral i -> "int"
  | BoolLiteral v -> "int"
  | Id v -> 
		if NameMap.mem v globals 
			then NameMap.find v globals
		else if NameMap.mem v locals
			then NameMap.find v locals
		else raise (Failure ("undeclared identifier " ^ v))
  | Not v -> ""
  | Binop (e1, op, e2) -> "int"
  | Assign (varName, e) -> "assign"
	| Call ("$print", e) -> "print" 
	| Call ("$println", e) -> "println"
	| Call ("$scan", e) -> "scan" 
	| Call ("$draw", e) -> "string" 
	| Call ("$value", e) ->"int" 
	| Call ("$deck", e) -> "int"
	| Call ("$label", e) ->"label" 
	| Call ("$break", e) ->"break" 
	| Call ("$getName", e) ->"string" 
	| Call ("$getValue", e) ->"int"
	| Call ("$getList", e) -> "element"
  | Call (f, actuals) -> 
		let func = NameMap.find f func_decls in func.freturn 
  | Noexpr -> ""
  | StringLiteral e -> "string"
	| LabelLiteral e -> "label"
	| ListLiteral e -> "list"
	| CardLiteral e -> "card"

let init_var =
  function 
		| "string" -> "\"\"" 
		| "int" -> "0" 
		| "bool" -> "1" 
		| "card" ->  "new Card()"
		| "list" -> "()"
		| x -> ""

let split_char sep str =
  let string_index_from i =
    try Some (String.index_from str i sep)
    with Not_found -> None
  in
  let rec aux i acc = match string_index_from i with
    | Some i' ->
        let w = String.sub str i (i' - i) in
        aux (succ i') (w::acc)
    | None ->
        let w = String.sub str i (String.length str - i) in
        List.rev (w::acc)
  in
  aux 0 []

(* Generates a string representation of variable decl *)
let string_of_var_decl var_decl =
	if var_decl.vartype = "list" 
	then "@" ^ (get_name var_decl.varname) ^ (" = " ^ ((init_var var_decl.vartype) ^ ";\n"))
  else var_decl.varname ^ (" = " ^ ((init_var var_decl.vartype) ^ ";\n"))
  
(* Generates a string representation of an expression*)
let rec string_of_expr locals globals func_decls =
  function
  | IntLiteral i -> string_of_int i
  | StringLiteral e -> e
  | Noexpr -> ""
  | Id var ->
      if NameMap.mem var locals
      then var
      else
        if NameMap.mem var globals
        then var
        else raise (Failure ("undeclared identifier " ^ var))
  | Not v -> "!" ^ (string_of_expr locals globals func_decls v)
  | BoolLiteral v -> (match v with | true -> "1" | false -> "0")
  | Binop (e1, op, e2) ->
		  let e1_data_type = get_data_type locals globals func_decls e1 in
			let e2_data_type = get_data_type locals globals func_decls e2 in
      let v1 = string_of_expr locals globals func_decls e1 in
      let v2 = string_of_expr locals globals func_decls e2 in
				(* type checking to make sure they are ints *)
				if ((e1_data_type <> "int") && (e1_data_type <> "bool")) then raise (Failure (v1 ^ " has to be of type int "))
				else if ((e2_data_type <> "int") && (e2_data_type <> "bool")) then raise (Failure (v2 ^ " has to be of type int "))
				else 
      let v3 =
        (match op with
         | Add -> v1 ^ ("+" ^ v2)
         | Sub -> v1 ^ ("-" ^ v2)
         | Mult -> v1 ^ ("*" ^ v2)
         | Div -> v1 ^ ("/" ^ v2)
         | Equal -> v1 ^ ("==" ^ v2)
         | Neq -> v1 ^ ("!=" ^ v2)
         | Less -> v1 ^ ("<" ^ v2)
         | Leq -> v1 ^ ("<=" ^ v2)
         | Greater -> v1 ^ (">" ^ v2)
         | Geq -> v1 ^ (">=" ^ v2)
         | Mod -> v1 ^ ("%" ^ v2)
         | Or -> v1 ^ ("||" ^ v2)
         | And -> v1 ^ ("&&" ^ v2))
      in "(" ^ (v3 ^ ")")
  | Assign (varName, e) ->
		  let e1_data_type = get_data_type locals globals func_decls varName in
			let e2_data_type = get_data_type locals globals func_decls e in
      let v1 = string_of_expr locals globals func_decls varName in
      let v2 = string_of_expr locals globals func_decls e in 
			(* Check to make sure what is being assigned is the correct datatype *)
			if (e2_data_type = "element" or e2_data_type = "scan")
			then (if NameMap.mem v1 locals or NameMap.mem v1 globals
        then (v1 ^ (" = " ^ v2))
				else raise (Failure ("undeclared identifier " ^ v1)))
			else (
				if (e1_data_type = "bool" && e2_data_type = "int") 
				then 
					( if NameMap.mem v1 locals or NameMap.mem v1 globals
	        then 
						(v1 ^ (" = " ^ v2))
					else raise (Failure ("undeclared identifier " ^ v1))) 
				else (
			if (e1_data_type <> e2_data_type) 
				then raise (Failure ("incompatible datatype during assignment. Expecting " ^ e1_data_type ^ " but " ^ e2_data_type ^ " is found" ))
			else
        if NameMap.mem v1 locals or NameMap.mem v1 globals
        then 
					(if (e1_data_type = "card") 
             then v1 ^ "->setname(\""^v2^"\")"
             else if (e1_data_type = "list")
						 	then ( "@" ^ (get_name v1) ^ " = " ^ v2) 
							else (v1 ^ (" = " ^ v2)))
				else raise (Failure ("undeclared identifier " ^ v1))))
	| LabelLiteral e -> ""
	| ListLiteral e ->  "(" ^(String.concat "," (List.map (string_of_expr locals globals func_decls) e)) ^ ")"
	| CardLiteral e -> e
  | Call ("$print", e) ->
      "print " ^
        (String.concat " "
           (List.map (string_of_expr locals globals func_decls) e))
  | Call ("$println", e) ->
      "print " ^
        ((String.concat " "
            (List.map (string_of_expr locals globals func_decls) e))
           ^ ".\"\\n\"")
  | Call ("$scan", e) ->
      let v1 =
        String.concat " "
          (List.map (string_of_expr locals globals func_decls) e)
      in
        if (String.length v1) > 0
        then raise (Failure "scan function does not take any arguments. ")
        else "<STDIN>"
	| Call ("$draw", e) ->
		(* wait until we add cardstacks to check . *)
			"$deck->draw()->getname()"
	| Call ("$printDeck", e) ->
		(* wait until we add cardstacks to check . *)
			"$deck->print()"	
	| Call ("$shuffleDeck", e) ->
		  "$deck->shuffle()"	
	| Call ("$value", e) ->
				(* check that v1 is a string that is valid. *)
			let v1 =
        String.concat " "
          (List.map (string_of_expr locals globals func_decls) e) 
			in
				"value("^v1^")"
	| Call ("$deck", e) ->
		(* function to see if the deck is nonempty *)
			 let v1 =
        String.concat " "
          (List.map (string_of_expr locals globals func_decls) e)
      in
        if (String.length v1) > 0
        then raise (Failure "deck function does not any arguments. ")
        else "$deck->size()"
	| Call ("$getName", e) -> 
		let v1 =
        String.concat " "
          (List.map (string_of_expr locals globals func_decls) e)
      in
        v1 ^ "-> getname()"
	| Call ("$getValue", e) -> 
				let v1 =
        (String.concat " "
          (List.map (string_of_expr locals globals func_decls) e))
      in
        v1 ^ "->getvalue()"
	| Call ("$getList", e) ->
			let l1 = (String.concat " " (List.map (string_of_expr locals globals func_decls) e)) in
			let l2 = (split_char ' ' l1) in
			let v1 = "$" ^ (String.sub (List.hd l2) 1 ((String.length  (List.hd l2))  -1)) in
			let v2 = (String.sub  (List.hd (List.tl l2)) 0 ((String.length (List.hd l2)) -1)) in 
       v1 ^ "[" ^ v2 ^ "]"		
	| Call ("$printList", e) ->
				let v1 =
        (String.concat " "
          (List.map (string_of_expr locals globals func_decls) e))
					in 
					let v2 = String.sub v1 1 ((String.length v1)-2) in
					"print @"^v2^";\nprint \"\\n\""
	| Call ("$sizeList", e) ->
				let v1 =
        (String.concat " "
          (List.map (string_of_expr locals globals func_decls) e))
					in 
				let v2 = String.sub v1 1 ((String.length v1)-2) in
					"scalar(@"^v2^")"
	| Call ("$push", e) ->
			let l1 = (String.concat " " (List.map (string_of_expr locals globals func_decls) e)) in
			let l2 = (split_char ' ' l1) in
			let v1 = "@" ^ (String.sub (List.hd l2) 1 ((String.length  (List.hd l2))  -1)) in
			let v2 = (String.sub (List.hd (List.tl l2)) 0 ((String.length  (List.hd l2))-1))in 
      "push(" ^ v1 ^ ",\"" ^ v2 ^ "\")"		
	| Call ("$pushInt", e) ->
			let l1 = (String.concat " " (List.map (string_of_expr locals globals func_decls) e)) in
			let l2 = (split_char ' ' l1) in
			let v1 = "@" ^ (String.sub (List.hd l2) 1 ((String.length  (List.hd l2))  -1)) in
			let v2 = (String.sub (List.hd (List.tl l2)) 0 ((String.length  (List.hd l2))-1))in 
      "push(" ^ v1 ^ "," ^ v2 ^ "\")"		
	| Call ("$label", e) ->
		(* label function for breaks. Check to see e is a string. *)
			let v1 =
        String.concat " "
          (List.map (string_of_expr locals globals func_decls) e)
			in
			 v1^":"
	| Call ("$break", e) ->
      (* should check that e is a string. Check that e is on list of labels. *)
			let v1 =
        String.concat " "
          (List.map (string_of_expr locals globals func_decls) e)
      in
        "last "^v1
  | Call (f, actuals) -> (* See if function is defined *)
      (* Check if function is defined *)
			if NameMap.mem f func_decls
      then let func = NameMap.find f func_decls in
				(* Check if the the number of arguments passed in matches expected num of args *)
				(* List.iter2 will give error if # of args are not accurate - Fatal error: exception Invalid_argument("List.iter2") *)
				(* fun x y will check if the data type matches *)
				List.iter2 (fun x y -> 
					if x.vartype <> get_data_type locals globals func_decls y
					then raise (Failure ("The arguments passed to function " ^ (get_name f) ^ " are type mismatched"))) 
						func.formals actuals ;
				(*(List.length actuals = List.length func.formals) then*)
        (get_name f) ^
          ("(" ^
             ((String.concat ", "
                 (List.map (string_of_expr locals globals func_decls) actuals))
                ^ ")"))
				(*else raise (Failure ("The number of arguments passed to function " ^ (get_name f) ^ " is mismatched"))*)
      else raise (Failure ("undefined function " ^ (get_name f)))
 
(* Generates a string representation of a statement*)
let rec string_of_stmt locals globals func_decls curr_func stmt =
  match stmt with
  | Block stmts ->
      "{\n" ^
        ((String.concat ""
            (List.map (fun x -> string_of_stmt locals globals func_decls curr_func x)
               stmts))
           ^ "}\n")
  | Expr e -> (string_of_expr locals globals func_decls e) ^ ";\n"
  | If (e, s, (Block [])) ->
		let data_type = get_data_type locals globals func_decls e in
		(*type checking *)
		if (data_type = "binop" || data_type = "bool" || data_type = "int") then
      "if (" ^
        ((string_of_expr locals globals func_decls e) ^
           (")\n {" ^ ((string_of_stmt globals locals func_decls curr_func s) ^ "}\n")))
		else raise (Failure ("if statement must take in a boolean expression "))
  | If (e, s1, s2) ->
		let data_type = get_data_type locals globals func_decls e in
	  (*type checking *)
		if (data_type = "binop" || data_type = "bool" || data_type = "int")
		then
      "if (" ^
        ((string_of_expr locals globals func_decls e) ^
           ("){" ^
              ((string_of_stmt locals globals func_decls curr_func s1) ^
                 ("}\n else {" ^
                    ((string_of_stmt locals globals func_decls curr_func s2) ^ "}\n")))))
		else raise (Failure ("if statement must take in a boolean expression "))
  | While (e, s) ->
		let data_type = get_data_type locals globals func_decls e in
	  (*type checking *)
		if (data_type = "binop" || data_type = "bool" || data_type = "int")
		then
      "while (" ^
        ((string_of_expr locals globals func_decls e) ^
           (") " ^ ((string_of_stmt locals globals func_decls curr_func s) ^ "\n")))
		else raise (Failure ("while statement must take in a boolean expression "))
  | For (e1, e2, e3, s) ->
		let data_type = get_data_type locals globals func_decls e2 in
		let e2_str = (string_of_expr locals globals func_decls e2) in
	  (*type checking *)
		if ((data_type = "binop" || data_type = "bool" || data_type = "int")|| (String.length e2_str)=0)
		then
      " for(" ^
        ((string_of_expr locals globals func_decls e1) ^
           ("; " ^
              (e2_str ^
                 ("; " ^
                    ((string_of_expr locals globals func_decls e3) ^
                       ("){ " ^
                          ((string_of_stmt locals globals func_decls curr_func s) ^
                             "} \n")))))))
		else raise (Failure ("for statement must take in a boolean expression as a second parameter"))
  | Return e -> 
		(* Checks if what is being returned match with what the function expects *)
		if ((get_data_type locals globals func_decls e) = curr_func.freturn) 
    then  "return " ^ ((string_of_expr locals globals func_decls e) ^ ";\n")
		else raise (Failure ("Function: " ^ (get_name curr_func.fname) ^ " needs to return of type: " ^ curr_func.freturn ))
	
	
let comp_func_main globals func_decls fdecl =
  (* Storing the function arguments in local symbol table*)
  (* checks if a argument name is used twice *)
  let locals =
    List.fold_left
      (fun lvars fargs ->
         if NameMap.mem fargs.varname lvars
          then
           raise
             (Failure (" Duplicate variable: " ^ (get_name fargs.varname)))
         else (* check to see if argument name is a keyword *)
           if is_keyword (get_name fargs.varname)
           then
             raise
               (Failure
                  ("variable name: " ^
                     ((get_name fargs.varname) ^ ", cannot be a keyword")))
           else NameMap.add fargs.varname fargs.vartype lvars)
      locals fdecl.formals in
  (* Storing the function local variables in local symbol table*)
  (* checks if a variable name is used twice *)
  let locals =
    List.fold_left
      (fun lvars fvars ->
         if NameMap.mem fvars.varname lvars
           then
           raise
             (Failure (" Duplicate variable: " ^ (get_name fvars.varname)))
         else (* check to see if argument name is a keyword *)
           if is_keyword (get_name fvars.varname)
           then
             raise
               (Failure
                  ("variable name: " ^
                     ((get_name fvars.varname) ^ ", cannot be a keyword")))
           else NameMap.add fvars.varname fvars.vartype lvars)
      locals fdecl.locals in
  (* Storing the string version of the function*)
  let func_str =
    String.concat "\n"
      (List.map (fun x -> string_of_stmt locals globals func_decls fdecl x)
         fdecl.body) in
  (* string representation of local *)
  let local_var_str =
    String.concat "\n" (List.map string_of_var_decl fdecl.locals)
  in
    (* return the string rep of the function *)
    "sub " ^
      ((get_name fdecl.fname) ^
         ("()\n{\n" ^ (local_var_str ^ (func_str ^ "}"))))
  
(* compiles a function *)
(* returns a string representation of the function *)
let rec comp_func globals func_decls fdecl =
  (* Storing the function arguments in local symbol table*)
  (* checks if a argument name is used twice *)
  let locals =
    List.fold_left
      (fun lvars fargs ->
         if NameMap.mem fargs.varname lvars
             then
           raise
             (Failure (" Duplicate variable: " ^ (get_name fargs.varname)))
         else (* check to see if argument name is a keyword *)
           if is_keyword (get_name fargs.varname)
           then
             raise
               (Failure
                  ("variable name: " ^
                     ((get_name fargs.varname) ^ ", cannot be a keyword")))
           else NameMap.add fargs.varname fargs.vartype lvars)
      locals fdecl.formals in
  (* Storing the function local variables in local symbol table*)
  (* checks if a variable name is used twice *)
  let locals =
    List.fold_left
      (fun lvars fvars ->
         if NameMap.mem fvars.varname lvars
           then
           raise
             (Failure (" Duplicate variable: " ^ (get_name fvars.varname)))
         else (* check to see if argument name is a keyword *)
           if is_keyword (get_name fvars.varname)
           then
             raise
               (Failure
                  ("variable name: " ^
                     ((get_name fvars.varname) ^ ", cannot be a keyword")))
           else NameMap.add fvars.varname fvars.vartype lvars)
      locals fdecl.locals in
  (* Storing the string version of the function*)
  let func_str =
    String.concat "\n"
      (List.map (fun x -> string_of_stmt locals globals func_decls fdecl x)
         fdecl.body) in
  (* string representation of local *)
  let local_var_str =
    String.concat "\n" (List.map string_of_var_decl fdecl.locals)
  in
    (* return the string rep of the function *)
    "sub " ^
      ((get_name fdecl.fname) ^
         ("()\n{\n" ^ (local_var_str ^ (func_str ^ "}"))))
  
(* compiles program *)
(* main entry point *)
(* returns a perl program *)
let compile (vars, funcs) = (* Put function declarations in a symbol table *)
 (* first function must be the main method *)
  let first_func = List.hd (List.rev funcs)
  in
    if first_func.fname <> "$main"
    then raise (Failure "The first function must be main")
    else
      (let func_decls =
         List.fold_left
           (fun funcs fdecl ->
              if NameMap.mem fdecl.fname funcs
              then
                raise
                  (Failure
                     (" function name must be unique. function name used already: "
                        ^ (get_name fdecl.fname)))
              else (* check to see if function name is a keyword *)
                if is_keyword (get_name fdecl.fname)
                then
                  raise
                    (Failure
                       ("function name: " ^
                          ((get_name fdecl.fname) ^ ", cannot be a keyword")))
                else NameMap.add fdecl.fname fdecl funcs)
           NameMap.empty funcs
       in
         (* must have a main function in program *)
         if ( != ) (NameMap.mem "$main" func_decls) true
         then raise (Failure " There must be a main function: ")
         else (* Put global variable declarations in a symbol table *)
           (* Checks to make sure the same name doesn't get used more than once*)
           (let globals =
              List.fold_left
                (fun gvars vdecl ->
                   if NameMap.mem vdecl.varname gvars
                   then
                     raise
                       (Failure
                          (" Duplicate variable: " ^ (get_name vdecl.varname)))
                   else (* check to see if variable name is a keyword *)
                     if is_keyword (get_name vdecl.varname)
                     then
                       raise
                         (Failure
                            ("variable name: " ^
                               ((get_name vdecl.varname) ^
                                  ", cannot be a keyword")))
                     else NameMap.add vdecl.varname vdecl.vartype gvars)
                globals vars in
            (* Returns a perl program *)
            let value_sub =
              "sub value {\n$name = $_[0];\n$value = substr $name, 1;\n" ^
                ("if ($value eq \"A\")\n{\nreturn 1;\n}\nelse\n" ^
                   ("{if ($value eq \"J\")\n{\nreturn 11;\n}\nelse\n{\n" ^
                      ("if ($value eq \"Q\")\n{\nreturn 12;\n}\nelse\n{\n" ^
                         ("if ($value eq \"K\")\n{\nreturn 13;\n}\nelse\n{\n"
                            ^ "return $value;\n}\n}\n}\n}\n}\n")))) in
            let perl_start =
              "#!/usr/local/bin/perl\nuse warnings;\n" ^
                "use Card;\nuse DefaultDeck;\n\n$deck = new DefaultDeck();\n"
            in
              perl_start ^ (* Global vars *)
                ((String.concat "\n" (List.map string_of_var_decl vars)) ^

                   (* Main function *)
                   ((comp_func_main globals func_decls
                       (List.hd (List.rev funcs)))
                      ^ (* value subroutine *)
                      ("\n" ^
                         (value_sub ^
                            ("\n" ^ (* functions *)
                               (String.concat "\n"
                                  (List.map
                                     (fun x -> comp_func globals func_decls x)
                                     (List.tl (List.rev funcs))))))))))
             ^ "\nprint main().\"\\n\";\n")
  
let print (vars, funcs) = (* Put function declarations in a symbol table *)
  let globals =
    List.fold_left
      (fun gvars vdecl ->
         if NameMap.mem vdecl.varname globals
         then raise (Failure (" Duplicate variable: " ^ vdecl.varname))
         else NameMap.add vdecl.varname vdecl.vartype gvars)
      globals vars
  in NameMap.iter (fun x y -> print_string (y ^ (" " ^ (x ^ "\n")))) globals